// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2009 Thomas Capricelli <orzel@freehackers.org>
// Copyright (C) 2012 Desire Nuentsa <desire.nuentsa_wakam@inria.fr>
//
// The algorithm of this class initially comes from MINPACK whose original authors are:
// Copyright Jorge More - Argonne National Laboratory
// Copyright Burt Garbow - Argonne National Laboratory
// Copyright Ken Hillstrom - Argonne National Laboratory
//
// This Source Code Form is subject to the terms of the Minpack license
// (a BSD-like license) described in the campaigned CopyrightMINPACK.txt file.
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef EIGEN_LEVENBERGMARQUARDT_H
#define EIGEN_LEVENBERGMARQUARDT_H

namespace Eigen {
namespace LevenbergMarquardtSpace {
    enum Status
    {
        NotStarted = -2,
        Running = -1,
        ImproperInputParameters = 0,
        RelativeReductionTooSmall = 1,
        RelativeErrorTooSmall = 2,
        RelativeErrorAndReductionTooSmall = 3,
        CosinusTooSmall = 4,
        TooManyFunctionEvaluation = 5,
        FtolTooSmall = 6,
        XtolTooSmall = 7,
        GtolTooSmall = 8,
        UserAsked = 9
    };
}

template <typename _Scalar, int NX = Dynamic, int NY = Dynamic> struct DenseFunctor
{
    typedef _Scalar Scalar;
    enum
    {
        InputsAtCompileTime = NX,
        ValuesAtCompileTime = NY
    };
    typedef Matrix<Scalar, InputsAtCompileTime, 1> InputType;
    typedef Matrix<Scalar, ValuesAtCompileTime, 1> ValueType;
    typedef Matrix<Scalar, ValuesAtCompileTime, InputsAtCompileTime> JacobianType;
    typedef ColPivHouseholderQR<JacobianType> QRSolver;
    const int m_inputs, m_values;

    DenseFunctor() : m_inputs(InputsAtCompileTime), m_values(ValuesAtCompileTime) {}
    DenseFunctor(int inputs, int values) : m_inputs(inputs), m_values(values) {}

    int inputs() const { return m_inputs; }
    int values() const { return m_values; }

    //int operator()(const InputType &x, ValueType& fvec) { }
    // should be defined in derived classes

    //int df(const InputType &x, JacobianType& fjac) { }
    // should be defined in derived classes
};

template <typename _Scalar, typename _Index> struct SparseFunctor
{
    typedef _Scalar Scalar;
    typedef _Index Index;
    typedef Matrix<Scalar, Dynamic, 1> InputType;
    typedef Matrix<Scalar, Dynamic, 1> ValueType;
    typedef SparseMatrix<Scalar, ColMajor, Index> JacobianType;
    typedef SparseQR<JacobianType, COLAMDOrdering<int>> QRSolver;
    enum
    {
        InputsAtCompileTime = Dynamic,
        ValuesAtCompileTime = Dynamic
    };

    SparseFunctor(int inputs, int values) : m_inputs(inputs), m_values(values) {}

    int inputs() const { return m_inputs; }
    int values() const { return m_values; }

    const int m_inputs, m_values;
    //int operator()(const InputType &x, ValueType& fvec) { }
    // to be defined in the functor

    //int df(const InputType &x, JacobianType& fjac) { }
    // to be defined in the functor if no automatic differentiation
};
namespace internal {
    template <typename QRSolver, typename VectorType>
    void lmpar2(const QRSolver& qr,
                const VectorType& diag,
                const VectorType& qtb,
                typename VectorType::Scalar m_delta,
                typename VectorType::Scalar& par,
                VectorType& x);
}
/**
  * \ingroup NonLinearOptimization_Module
  * \brief Performs non linear optimization over a non-linear function,
  * using a variant of the Levenberg Marquardt algorithm.
  *
  * Check wikipedia for more information.
  * http://en.wikipedia.org/wiki/Levenberg%E2%80%93Marquardt_algorithm
  */
template <typename _FunctorType> class LevenbergMarquardt : internal::no_assignment_operator
{
public:
    typedef _FunctorType FunctorType;
    typedef typename FunctorType::QRSolver QRSolver;
    typedef typename FunctorType::JacobianType JacobianType;
    typedef typename JacobianType::Scalar Scalar;
    typedef typename JacobianType::RealScalar RealScalar;
    typedef typename QRSolver::StorageIndex PermIndex;
    typedef Matrix<Scalar, Dynamic, 1> FVectorType;
    typedef PermutationMatrix<Dynamic, Dynamic, int> PermutationType;

public:
    LevenbergMarquardt(FunctorType& functor) : m_functor(functor), m_nfev(0), m_njev(0), m_fnorm(0.0), m_gnorm(0), m_isInitialized(false), m_info(InvalidInput)
    {
        resetParameters();
        m_useExternalScaling = false;
    }

    LevenbergMarquardtSpace::Status minimize(FVectorType& x);
    LevenbergMarquardtSpace::Status minimizeInit(FVectorType& x);
    LevenbergMarquardtSpace::Status minimizeOneStep(FVectorType& x);
    LevenbergMarquardtSpace::Status lmder1(FVectorType& x, const Scalar tol = std::sqrt(NumTraits<Scalar>::epsilon()));
    static LevenbergMarquardtSpace::Status
    lmdif1(FunctorType& functor, FVectorType& x, Index* nfev, const Scalar tol = std::sqrt(NumTraits<Scalar>::epsilon()));

    /** Sets the default parameters */
    void resetParameters()
    {
        using std::sqrt;

        m_factor = 100.;
        m_maxfev = 400;
        m_ftol = sqrt(NumTraits<RealScalar>::epsilon());
        m_xtol = sqrt(NumTraits<RealScalar>::epsilon());
        m_gtol = 0.;
        m_epsfcn = 0.;
    }

    /** Sets the tolerance for the norm of the solution vector*/
    void setXtol(RealScalar xtol) { m_xtol = xtol; }

    /** Sets the tolerance for the norm of the vector function*/
    void setFtol(RealScalar ftol) { m_ftol = ftol; }

    /** Sets the tolerance for the norm of the gradient of the error vector*/
    void setGtol(RealScalar gtol) { m_gtol = gtol; }

    /** Sets the step bound for the diagonal shift */
    void setFactor(RealScalar factor) { m_factor = factor; }

    /** Sets the error precision  */
    void setEpsilon(RealScalar epsfcn) { m_epsfcn = epsfcn; }

    /** Sets the maximum number of function evaluation */
    void setMaxfev(Index maxfev) { m_maxfev = maxfev; }

    /** Use an external Scaling. If set to true, pass a nonzero diagonal to diag() */
    void setExternalScaling(bool value) { m_useExternalScaling = value; }

    /** \returns the tolerance for the norm of the solution vector */
    RealScalar xtol() const { return m_xtol; }

    /** \returns the tolerance for the norm of the vector function */
    RealScalar ftol() const { return m_ftol; }

    /** \returns the tolerance for the norm of the gradient of the error vector */
    RealScalar gtol() const { return m_gtol; }

    /** \returns the step bound for the diagonal shift */
    RealScalar factor() const { return m_factor; }

    /** \returns the error precision */
    RealScalar epsilon() const { return m_epsfcn; }

    /** \returns the maximum number of function evaluation */
    Index maxfev() const { return m_maxfev; }

    /** \returns a reference to the diagonal of the jacobian */
    FVectorType& diag() { return m_diag; }

    /** \returns the number of iterations performed */
    Index iterations() { return m_iter; }

    /** \returns the number of functions evaluation */
    Index nfev() { return m_nfev; }

    /** \returns the number of jacobian evaluation */
    Index njev() { return m_njev; }

    /** \returns the norm of current vector function */
    RealScalar fnorm() { return m_fnorm; }

    /** \returns the norm of the gradient of the error */
    RealScalar gnorm() { return m_gnorm; }

    /** \returns the LevenbergMarquardt parameter */
    RealScalar lm_param(void) { return m_par; }

    /** \returns a reference to the  current vector function 
     */
    FVectorType& fvec() { return m_fvec; }

    /** \returns a reference to the matrix where the current Jacobian matrix is stored
     */
    JacobianType& jacobian() { return m_fjac; }

    /** \returns a reference to the triangular matrix R from the QR of the jacobian matrix.
     * \sa jacobian()
     */
    JacobianType& matrixR() { return m_rfactor; }

    /** the permutation used in the QR factorization
     */
    PermutationType permutation() { return m_permutation; }

    /** 
     * \brief Reports whether the minimization was successful
     * \returns \c Success if the minimization was successful,
     *         \c NumericalIssue if a numerical problem arises during the 
     *          minimization process, for example during the QR factorization
     *         \c NoConvergence if the minimization did not converge after 
     *          the maximum number of function evaluation allowed
     *          \c InvalidInput if the input matrix is invalid
     */
    ComputationInfo info() const { return m_info; }

private:
    JacobianType m_fjac;
    JacobianType m_rfactor;  // The triangular matrix R from the QR of the jacobian matrix m_fjac
    FunctorType& m_functor;
    FVectorType m_fvec, m_qtf, m_diag;
    Index n;
    Index m;
    Index m_nfev;
    Index m_njev;
    RealScalar m_fnorm;   // Norm of the current vector function
    RealScalar m_gnorm;   //Norm of the gradient of the error
    RealScalar m_factor;  //
    Index m_maxfev;       // Maximum number of function evaluation
    RealScalar m_ftol;    //Tolerance in the norm of the vector function
    RealScalar m_xtol;    //
    RealScalar m_gtol;    //tolerance of the norm of the error gradient
    RealScalar m_epsfcn;  //
    Index m_iter;         // Number of iterations performed
    RealScalar m_delta;
    bool m_useExternalScaling;
    PermutationType m_permutation;
    FVectorType m_wa1, m_wa2, m_wa3, m_wa4;  //Temporary vectors
    RealScalar m_par;
    bool m_isInitialized;  // Check whether the minimization step has been called
    ComputationInfo m_info;
};

template <typename FunctorType> LevenbergMarquardtSpace::Status LevenbergMarquardt<FunctorType>::minimize(FVectorType& x)
{
    LevenbergMarquardtSpace::Status status = minimizeInit(x);
    if (status == LevenbergMarquardtSpace::ImproperInputParameters)
    {
        m_isInitialized = true;
        return status;
    }
    do
    {
        //       std::cout << " uv " << x.transpose() << "\n";
        status = minimizeOneStep(x);
    } while (status == LevenbergMarquardtSpace::Running);
    m_isInitialized = true;
    return status;
}

template <typename FunctorType> LevenbergMarquardtSpace::Status LevenbergMarquardt<FunctorType>::minimizeInit(FVectorType& x)
{
    n = x.size();
    m = m_functor.values();

    m_wa1.resize(n);
    m_wa2.resize(n);
    m_wa3.resize(n);
    m_wa4.resize(m);
    m_fvec.resize(m);
    //FIXME Sparse Case : Allocate space for the jacobian
    m_fjac.resize(m, n);
    //     m_fjac.reserve(VectorXi::Constant(n,5)); // FIXME Find a better alternative
    if (!m_useExternalScaling)
        m_diag.resize(n);
    eigen_assert((!m_useExternalScaling || m_diag.size() == n) && "When m_useExternalScaling is set, the caller must provide a valid 'm_diag'");
    m_qtf.resize(n);

    /* Function Body */
    m_nfev = 0;
    m_njev = 0;

    /*     check the input parameters for errors. */
    if (n <= 0 || m < n || m_ftol < 0. || m_xtol < 0. || m_gtol < 0. || m_maxfev <= 0 || m_factor <= 0.)
    {
        m_info = InvalidInput;
        return LevenbergMarquardtSpace::ImproperInputParameters;
    }

    if (m_useExternalScaling)
        for (Index j = 0; j < n; ++j)
            if (m_diag[j] <= 0.)
            {
                m_info = InvalidInput;
                return LevenbergMarquardtSpace::ImproperInputParameters;
            }

    /*     evaluate the function at the starting point */
    /*     and calculate its norm. */
    m_nfev = 1;
    if (m_functor(x, m_fvec) < 0)
        return LevenbergMarquardtSpace::UserAsked;
    m_fnorm = m_fvec.stableNorm();

    /*     initialize levenberg-marquardt parameter and iteration counter. */
    m_par = 0.;
    m_iter = 1;

    return LevenbergMarquardtSpace::NotStarted;
}

template <typename FunctorType> LevenbergMarquardtSpace::Status LevenbergMarquardt<FunctorType>::lmder1(FVectorType& x, const Scalar tol)
{
    n = x.size();
    m = m_functor.values();

    /* check the input parameters for errors. */
    if (n <= 0 || m < n || tol < 0.)
        return LevenbergMarquardtSpace::ImproperInputParameters;

    resetParameters();
    m_ftol = tol;
    m_xtol = tol;
    m_maxfev = 100 * (n + 1);

    return minimize(x);
}

template <typename FunctorType>
LevenbergMarquardtSpace::Status LevenbergMarquardt<FunctorType>::lmdif1(FunctorType& functor, FVectorType& x, Index* nfev, const Scalar tol)
{
    Index n = x.size();
    Index m = functor.values();

    /* check the input parameters for errors. */
    if (n <= 0 || m < n || tol < 0.)
        return LevenbergMarquardtSpace::ImproperInputParameters;

    NumericalDiff<FunctorType> numDiff(functor);
    // embedded LevenbergMarquardt
    LevenbergMarquardt<NumericalDiff<FunctorType>> lm(numDiff);
    lm.setFtol(tol);
    lm.setXtol(tol);
    lm.setMaxfev(200 * (n + 1));

    LevenbergMarquardtSpace::Status info = LevenbergMarquardtSpace::Status(lm.minimize(x));
    if (nfev)
        *nfev = lm.nfev();
    return info;
}

}  // end namespace Eigen

#endif  // EIGEN_LEVENBERGMARQUARDT_H
